# api/auth.py
from fastapi import APIRouter, HTTPException, Request, Response
from fastapi import File, UploadFile
from pydantic import BaseModel
from httpx import AsyncClient
from openai import OpenAI
from fastapi.responses import StreamingResponse, JSONResponse
import asyncio
from sse_starlette import EventSourceResponse
import json
import os

import fitz


def extract_text_from_pdf(pdf_path):
    document = fitz.open(pdf_path)
    text = ""

    for page_num in range(len(document)):
        page = document[page_num]
        # 获取页面的文本（使用fitz.TO_UNICODE_TEXT选项）
        page_text = page.get_text("text")
        text += page_text

    document.close()
    return text


test_router = APIRouter()


@test_router.get("/api/testget")
async def test_get(username: str = None, password: str = None):
    if username and password:
        # 这里执行登录逻辑
        return {"message": f"Hello, {username}!"}
    else:
        return {"message": "Missing username or password"}

# 定义请求数据模型


class TestLoginRequest(BaseModel):
    username: str
    password: str

# 定义一个POST路由，使用请求数据模型来验证请求体


@test_router.post("/api/testpost")
async def test_post(login_data: TestLoginRequest):
    # 这里执行登录逻辑
    if login_data.username and login_data.password:
        # 假设登录成功
        return {"message": f"Hello, {login_data.username}!"}
    else:
        # 如果请求数据不完整，抛出异常
        raise HTTPException(
            status_code=400, detail="Missing username or password")


@test_router.get("/api/teststream")
async def flush_stream(request: Request):
    async def event_generator(request: Request):
        res_str = "这是一个流式输出他会将每个字挨个挨个的输出哈哈！！！"
        for idx, word in enumerate(res_str):
            if await request.is_disconnected():
                print("连接已中断")
                break
            data = json.dumps({"id": idx, "chunk": word}, ensure_ascii=False)
            yield data
            await asyncio.sleep(1)

    return EventSourceResponse(event_generator(request))


MOONSHOT_API_KEY = 'sk-26HaIgGrRMDLdPku4ooFSiXNhvf0Jb0c35m6tFCJRYNMsCPx'  # from chen lei


@test_router.get("/api/testchat")
async def chat(request: Request):
    async def make_kimi_requests(prompt: str):
        client = OpenAI(
            api_key=MOONSHOT_API_KEY,
            base_url="https://api.moonshot.cn/v1",
        )

        response = client.chat.completions.create(
            model="moonshot-v1-8k",
            messages=[
                {"role": "system",
                 "content": '你是一个人工智能助手，你将回答用户的问题或者帮助用户完成任务'},
                {"role": "user",
                 "content": prompt}
            ],
            temperature=0.7,
            top_p=1,
            stream=True,
        )
        for idx, chunk in enumerate(response):
            if await request.is_disconnected():
                print("[api test]error: 连接已中断")
                break
            chunk_message = chunk.choices[0].delta
            if not chunk_message.content:
                continue
            data = json.dumps(
                {"id": idx, "chunk": chunk_message.content}, ensure_ascii=False)
            yield data
            await asyncio.sleep(0.5)  # 模拟卡顿，充VIP解锁
    prompt = "请为我写一段话，100字，描述朱自清的荷塘月色"
    return EventSourceResponse(make_kimi_requests(prompt))


@test_router.post("/api/uploadPdf")
async def upload_pdf(file: UploadFile = File(...)):
    data_directory = "./data/pdf/"
    if not os.path.exists(data_directory):
        os.makedirs(data_directory)
    # try:
    filename = file.filename

    if not filename.endswith(".pdf"):
        return JSONResponse(status_code=400, content={"message": "Only PDF files are allowed"})

    save_path = os.path.join(data_directory, filename)
    with open(save_path, "wb") as f:
        contents = await file.read()
        f.write(contents)

    extracted_text = extract_text_from_pdf(save_path)

    return JSONResponse(status_code=200, content={"message": "File uploaded successfully", "filename": filename, "extracted_text": extracted_text})

    # except Exception as e:
    #     return JSONResponse(status_code=500, content={"message": str(e)})
